"""
CBMPy: Constraint Based Modelling in Python (http://pysces.sourceforge.net/cbm)
============
Copyright (C) 2009-2012 Brett G. Olivier, VU University Amsterdam, Amsterdam, The Netherlands

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>

Author: Brett G. Olivier
Contact email: bgoli@users.sourceforge.net
Last edit: $Author: bgoli $ ($Id: CBNetDB.py 990 2012-02-15 13:52:11Z bgoli $)

"""

import os, time, numpy

class DBTools:
    sqlite = None
    db_conn = None
    db_cursor = None
    
    def __init__(self):
        from pysqlite2 import dbapi2 as sqlite
        self.sqlite = sqlite

    def connectSQLiteDB(self, db_name, work_dir):
        # connect to Ecoli DB
        self.db_conn = self.sqlite.connect(os.path.join(work_dir, db_name))
        self.db_cursor = self.db_conn.cursor()

    def createDBTable(self, table, sqlcols=['gene TEXT PRIMARY KEY', 'aa_seq TEXT', 'nuc_seq TEXT', 'aa_len INT', 'nuc_len INT']):
        "CREATE TABLE %s (gene TEXT PRIMARY KEY, aa_seq TEXT, nuc_seq TEXT, aa_len INT, nuc_len INT) % table"

        SQL = 'CREATE TABLE %s (' % table
        for c in sqlcols:
            SQL += ' %s,' % c
        SQL = SQL[:-1]
        SQL += ' )'
        print SQL
        try:
            self.db_cursor.execute('SELECT * FROM %s' % table)
            print 'Table %s exists' % table
        except self.sqlite.OperationalError:
            print 'Table %s does not exist, creating it' % table
            self.db_cursor.execute(SQL)

    def fetchAll(self, sql):
        """E.g. SELECT aa_len FROM ecoli_genes WHERE gene="G"'"""
        print sql
        r = None
        try:
            r = self.db_cursor.execute(sql).fetchall()
        except Exception, ex:
            print ex
        return r

class KeGGTools:
    Kclient = None
    
    def __init__(self, url):
        import suds
        self.Kclient = suds.client.Client(url)

    def fetchSeqfromKeGG(self, k_gene):
        g2 = 'None'
        p2 = 'None'
        try:
            g = self.Kclient.service.bget("-f -n n %s" % k_gene)
            if g == None:
                print '\n*****\nWARNING: potential naming error in gene: %s!!\n*****\n' % k_gene
            g2 = g.split('(N)')[1].replace('\n','')
        except Exception, ex:
            print '\nGene sequence get exception (%s)!\n' % k_gene
            print ex
        try:
            p = self.Kclient.service.bget("-f -n a %s" % k_gene)
            if p == None:
                print '\n*****\nWARNING: potential naming error in gene: %s!!\n*****\n' % k_gene
            p2 = p.split('(A)')[1].replace('\n','')
        except Exception, ex:
            print '\nProtein sequence get exception (%s)!\n' % k_gene
            print ex
        return g2, p2


class KeGGSequenceTools:
    DB = None
    KEGG = None
    
    def __init__(self, url, db_name, work_dir):
        self.DB = DBTools()
        self.DB.connectSQLiteDB(db_name, work_dir)
        self.KEGG = KeGGTools(url)
    
    def buildGeneDatabase(self, genes, tablename, UPDATE_IF_EXISTS=False, default_length=0):
        cntr = 1
        cntr2 = 1

        for ecg in genes:
            print 'Processing gene %s of %s' % (cntr, len(genes))
            entry_exists = False
            testg = self.DB.db_cursor.execute('SELECT * FROM %s WHERE gene="%s" ' % (tablename,ecg)).fetchall()
            if len(testg) > 0:
                entry_exists = True

            ##  if ecg in ['eco:b1898','eco:b1899','eco:b3692','eco:b3111','eco:b4228','eco:b2978',\
                ##  'eco:b1416','eco:b3112','eco:b1417','eco:b3768','eco:b3767','eco:b4229']:
                ##  raw_input(testg)

            tstart = time.time()
            if not entry_exists:
                print '\tadding gene %s' % ecg,
                gene2, prot2 = self.KEGG.fetchSeqfromKeGG(ecg)
                if gene2 != 'None' and prot2 != 'None':
                    self.DB.db_cursor.execute("INSERT INTO %s (gene, aa_seq, nuc_seq, aa_len, nuc_len) VALUES (?, ?, ?, ?, ?)" % tablename,
                                    (str(ecg), str(prot2), str(gene2), int(len(prot2)), int(len(gene2))) )
                else:
                    print '\nGene %s cannot be found and is probably an incorrect annotation assigning length: %s\n' % (ecg, default_length)
                    self.DB.db_cursor.execute("INSERT INTO %s (gene, aa_seq, nuc_seq, aa_len, nuc_len) VALUES (?, ?, ?, ?, ?)" % tablename,
                                    (str(ecg), 'None', 'None', default_length, default_length) )
            elif entry_exists and UPDATE_IF_EXISTS:
                print '\tupdating gene %s' % ecg,
                gene2, prot2 = self.fetchSeqfromKeGG(ecg)
                self.DB.db_cursor.execute('UPDATE %s SET aa_seq="%s", nuc_seq="%s", aa_len="%s", nuc_len="%s" WHERE gene="%s"' % (tablename, prot2, gene2, int(len(prot2)), int(len(gene2)), ecg))
            else:
                print '\tskipping gene %s' % ecg,
            tend = time.time()
            print ' ... done (%2.2f).' % (tend-tstart)
            #  if cntr == 6:
                #  break
            cntr += 1
            cntr2 += 1
            if cntr2 == 21:
                self.DB.db_cursor.connection.commit()
                cntr2 = 1
                
    def getPeptideLengthsFromDB(self, genes, keg_prefix):
        gene_peplen = {}
        for G in genes:
            print G
            ##  Glen = self.cursor.execute('SELECT aa_len FROM ecoli_genes WHERE gene="%s"' % G).fetchall()[0][0]
            Glen = self.DB.fetchAll('SELECT aa_len FROM ecoli_genes WHERE gene="%s"' % G)[0][0]
            print Glen
            gene_peplen.update({G.replace(keg_prefix,'') : Glen})
        return gene_peplen
